{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "883a7eba",
   "metadata": {},
   "source": [
    "# Daily Dose of Data Science\n",
    "\n",
    "This notebook accompanies the code for our model compression blog.\n",
    "\n",
    "Read the full blog here: [Machine Learning Model Compression: A Critical Step Towards Efficient Deep Learning](https://www.dailydoseofds.com/model-compression-a-critical-step-towards-efficient-machine-learning)\n",
    "\n",
    "Author: Avi Chawla"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62eba7f5",
   "metadata": {},
   "source": [
    "# Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4875592c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from time import time\n",
    "from tqdm import tqdm\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b526182",
   "metadata": {},
   "source": [
    "# Load the MNIST dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "730908bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])\n",
    "trainset = torchvision.datasets.MNIST(root='./data', train=True, download=True, transform=transform)\n",
    "trainloader = DataLoader(trainset, batch_size=64, shuffle=True)\n",
    "\n",
    "testset = torchvision.datasets.MNIST(root='./data', train=False, download=True, transform=transform)\n",
    "testloader = DataLoader(testset, batch_size=64, shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72dae106",
   "metadata": {},
   "source": [
    "# Knowledge Distillation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acd36b45",
   "metadata": {},
   "source": [
    "## Teacher Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "30903925",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TeacherNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(TeacherNet, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, 32, 5)\n",
    "        self.pool = nn.MaxPool2d(5, 5)\n",
    "        self.fc1 = nn.Linear(32 * 4 * 4, 128)\n",
    "        self.fc2 = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = F.relu(self.conv1(x))\n",
    "        x = self.pool(x)\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = self.fc2(x)\n",
    "        return x        "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ccc08e55",
   "metadata": {},
   "source": [
    "## Evaluation function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "0031cddd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model):\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    with torch.no_grad():\n",
    "        for data in testloader:\n",
    "            inputs, labels = data\n",
    "            outputs = model(inputs)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "    return correct / total"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84085c9c",
   "metadata": {},
   "source": [
    "## Initialize and train the teacher model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "edd90fe2",
   "metadata": {},
   "outputs": [],
   "source": [
    "teacher_model = TeacherNet()\n",
    "teacher_optimizer = optim.Adam(teacher_model.parameters(), lr=0.001)\n",
    "teacher_criterion = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "063145bc",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Loss: 0.23366064861265898, Accuracy: 97.60%\n",
      "Epoch 2, Loss: 0.07699692965661889, Accuracy: 98.00%\n",
      "Epoch 3, Loss: 0.058064278137973394, Accuracy: 98.44%\n",
      "Epoch 4, Loss: 0.04937064894107677, Accuracy: 98.24%\n",
      "Epoch 5, Loss: 0.04162352114517703, Accuracy: 98.53%\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(5):\n",
    "    teacher_model.train()\n",
    "    running_loss = 0.0\n",
    "    \n",
    "    for data in trainloader:\n",
    "        inputs, labels = data\n",
    "        teacher_optimizer.zero_grad()\n",
    "        outputs = teacher_model(inputs)\n",
    "        loss = teacher_criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        teacher_optimizer.step()\n",
    "        \n",
    "        running_loss += loss.item()\n",
    "        \n",
    "    teacher_accuracy = evaluate(teacher_model)\n",
    "        \n",
    "    print(f\"Epoch {epoch + 1}, Loss: {running_loss / len(trainloader)}, Accuracy: {teacher_accuracy * 100:.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c0d1678",
   "metadata": {},
   "source": [
    "## Student Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e62c6192",
   "metadata": {},
   "outputs": [],
   "source": [
    "class StudentNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(StudentNet, self).__init__()\n",
    "        self.fc1 = nn.Linear(28 * 28, 128)\n",
    "        self.fc2 = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = F.relu(self.fc1(x))\n",
    "        x = self.fc2(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57741f49",
   "metadata": {},
   "source": [
    "## Initialize and train the teacher model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3c3291f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "student_model = StudentNet()\n",
    "student_optimizer = optim.Adam(student_model.parameters(), lr=0.001)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ca726c7",
   "metadata": {},
   "source": [
    "## Loss function (KL Divergence)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "8dedb872",
   "metadata": {},
   "outputs": [],
   "source": [
    "def knowledge_distillation_loss(student_logits, teacher_logits):\n",
    "    p_teacher = F.softmax(teacher_logits , dim=1)\n",
    "    p_student = F.log_softmax(student_logits, dim=1)\n",
    "    loss = F.kl_div(p_student, p_teacher, reduction='batchmean')\n",
    "    return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5b075f35",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Loss: 1.97617478094473, Accuracy: 93.53%\n",
      "Epoch 2, Loss: 0.9071605966373044, Accuracy: 94.67%\n",
      "Epoch 3, Loss: 0.6211776698874251, Accuracy: 96.30%\n",
      "Epoch 4, Loss: 0.48355193005483244, Accuracy: 96.29%\n",
      "Epoch 5, Loss: 0.4033386060778218, Accuracy: 96.34%\n"
     ]
    }
   ],
   "source": [
    "# Train the student model with knowledge distillation\n",
    "for epoch in range(5):  # You can adjust the number of epochs\n",
    "    student_model.train()\n",
    "    running_loss = 0.0\n",
    "    \n",
    "    for data in trainloader:\n",
    "        inputs, labels = data\n",
    "        student_optimizer.zero_grad()\n",
    "        student_logits = student_model(inputs)\n",
    "        teacher_logits = teacher_model(inputs).detach()  # Detach the teacher's output to avoid backpropagation\n",
    "        loss = knowledge_distillation_loss(student_logits, teacher_logits)\n",
    "        loss.backward()\n",
    "        student_optimizer.step()\n",
    "        \n",
    "        running_loss += loss.item()\n",
    "    \n",
    "    student_accuracy = evaluate(student_model)\n",
    "        \n",
    "    print(f\"Epoch {epoch + 1}, Loss: {running_loss / len(testloader)}, Accuracy: {student_accuracy * 100:.2f}%\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "4b2c39e1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.61 s ± 21.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit evaluate(teacher_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b36d61f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.09 s ± 63 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n"
     ]
    }
   ],
   "source": [
    "%timeit evaluate(student_model) # student model runs faster"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ee9a169",
   "metadata": {},
   "source": [
    "# Zero-Pruning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82037888",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "79249e2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a simple neural network for MNIST classification\n",
    "class SimpleNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(SimpleNet, self).__init__()\n",
    "        self.fc1 = nn.Linear(28 * 28, 512)\n",
    "        self.fc2 = nn.Linear(512, 256)\n",
    "        self.fc3 = nn.Linear(256, 128)\n",
    "        self.fc4 = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x = torch.relu(self.fc1(x))\n",
    "        x = torch.relu(self.fc2(x))\n",
    "        x = torch.relu(self.fc3(x))\n",
    "        x = self.fc4(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "397126d3",
   "metadata": {},
   "source": [
    "## Evaluation function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "34af1ff2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model):\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    with torch.no_grad():\n",
    "        for data in testloader:\n",
    "            inputs, labels = data\n",
    "            outputs = model(inputs)\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "    return correct / total"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27aaab20",
   "metadata": {},
   "source": [
    "## Initialize and train the neural network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "29530d8e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Loss: 0.0002394356866126884, Accuracy: 95.17%\n",
      "Epoch 2, Loss: 0.0002629532933489346, Accuracy: 96.14%\n",
      "Epoch 3, Loss: 5.20410822398627e-05, Accuracy: 96.73%\n",
      "Epoch 4, Loss: 5.284582437482724e-05, Accuracy: 97.09%\n",
      "Epoch 5, Loss: 2.2851257845123947e-06, Accuracy: 97.07%\n"
     ]
    }
   ],
   "source": [
    "net = SimpleNet()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(net.parameters(), lr=0.001)\n",
    "\n",
    "for epoch in range(5):\n",
    "    running_loss = 0.0\n",
    "    net.train()\n",
    "    for data in trainloader:\n",
    "        inputs, labels = data\n",
    "        optimizer.zero_grad()\n",
    "        outputs = net(inputs.view(-1, 28*28))\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "    running_loss += loss.item()\n",
    "        \n",
    "    accuracy = evaluate(net)\n",
    "        \n",
    "    print(f\"Epoch {epoch + 1}, Loss: {running_loss / len(trainloader)}, Accuracy: {accuracy * 100:.2f}%\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "a1ad7553",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save model\n",
    "torch.save(net, \"net.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f805a8d",
   "metadata": {},
   "source": [
    "## Define the pruning threshold (λ) and apply Pruning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "28450360",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████| 11/11 [00:14<00:00,  1.33s/it]\n"
     ]
    }
   ],
   "source": [
    "thresholds = np.linspace(0, 0.1, 11) # range of thresholds\n",
    "results = []\n",
    "total_params = np.sum([param.numel() for name, param in net.named_parameters() if 'weight' in name])\n",
    "\n",
    "for threshold in tqdm(thresholds):\n",
    "    \n",
    "    # apply zero pruning\n",
    "    for name, param in net.named_parameters():\n",
    "        if 'weight' in name: # Apply pruning to weight parameters\n",
    "            param.data[torch.abs(param.data) < threshold] = 0\n",
    "            \n",
    "    # Count zero param\n",
    "    zero_params = np.sum([torch.sum(param == 0).item() for name, param in net.named_parameters() if 'weight' in name])\n",
    "    \n",
    "    accuracy = evaluate(net)\n",
    "    \n",
    "    results.append([threshold, accuracy, total_params, zero_params])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "699c465b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><div id=5bc217a7-d7e5-422c-8b5b-7391f58b4a17 style=\"display:none; background-color:#9D6CFF; color:white; width:200px; height:30px; padding-left:5px; border-radius:4px; flex-direction:row; justify-content:space-around; align-items:center;\" onmouseover=\"this.style.backgroundColor='#BA9BF8'\" onmouseout=\"this.style.backgroundColor='#9D6CFF'\" onclick=\"window.commands?.execute('create-mitosheet-from-dataframe-output');\">See Full Dataframe in Mito</div> <script> if (window.commands?.hasCommand('create-mitosheet-from-dataframe-output')) document.getElementById('5bc217a7-d7e5-422c-8b5b-7391f58b4a17').style.display = 'flex' </script> <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Threshold</th>\n",
       "      <th>Accuracy</th>\n",
       "      <th>Original Params</th>\n",
       "      <th>Zero Params</th>\n",
       "      <th>Zero percentage</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.00</td>\n",
       "      <td>0.9707</td>\n",
       "      <td>566528</td>\n",
       "      <td>0</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0.01</td>\n",
       "      <td>0.9705</td>\n",
       "      <td>566528</td>\n",
       "      <td>120825</td>\n",
       "      <td>21.327278</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0.02</td>\n",
       "      <td>0.9703</td>\n",
       "      <td>566528</td>\n",
       "      <td>238169</td>\n",
       "      <td>42.040111</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0.03</td>\n",
       "      <td>0.9658</td>\n",
       "      <td>566528</td>\n",
       "      <td>345597</td>\n",
       "      <td>61.002634</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0.04</td>\n",
       "      <td>0.9375</td>\n",
       "      <td>566528</td>\n",
       "      <td>422488</td>\n",
       "      <td>74.574955</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>0.06</td>\n",
       "      <td>0.8220</td>\n",
       "      <td>566528</td>\n",
       "      <td>488538</td>\n",
       "      <td>86.233690</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>0.07</td>\n",
       "      <td>0.7939</td>\n",
       "      <td>566528</td>\n",
       "      <td>505050</td>\n",
       "      <td>89.148286</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>0.08</td>\n",
       "      <td>0.7405</td>\n",
       "      <td>566528</td>\n",
       "      <td>517510</td>\n",
       "      <td>91.347647</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>0.09</td>\n",
       "      <td>0.6363</td>\n",
       "      <td>566528</td>\n",
       "      <td>527475</td>\n",
       "      <td>93.106607</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>0.10</td>\n",
       "      <td>0.5383</td>\n",
       "      <td>566528</td>\n",
       "      <td>535227</td>\n",
       "      <td>94.474942</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table></div>"
      ],
      "text/plain": [
       "    Threshold  Accuracy  Original Params  Zero Params  Zero percentage\n",
       "0        0.00    0.9707           566528            0         0.000000\n",
       "1        0.01    0.9705           566528       120825        21.327278\n",
       "2        0.02    0.9703           566528       238169        42.040111\n",
       "3        0.03    0.9658           566528       345597        61.002634\n",
       "4        0.04    0.9375           566528       422488        74.574955\n",
       "5        0.05    0.8780           566528       465046        82.087028\n",
       "6        0.06    0.8220           566528       488538        86.233690\n",
       "7        0.07    0.7939           566528       505050        89.148286\n",
       "8        0.08    0.7405           566528       517510        91.347647\n",
       "9        0.09    0.6363           566528       527475        93.106607\n",
       "10       0.10    0.5383           566528       535227        94.474942"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results = pd.DataFrame(results, columns = [\"Threshold\", \"Accuracy\", \"Original Params\", \"Zero Params\"])\n",
    "results[\"Zero percentage\"] = 100*results[\"Zero Params\"]/results[\"Original Params\"]\n",
    "results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3007ad6",
   "metadata": {},
   "source": [
    "## Select best threshold and apply pruning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "b5217795",
   "metadata": {},
   "outputs": [],
   "source": [
    "threshold = 0.03\n",
    "net = torch.load('net.pt')\n",
    "\n",
    "# apply zero pruning\n",
    "for name, param in net.named_parameters():\n",
    "    if 'weight' in name: # Apply pruning to weight parameters\n",
    "        param.data[torch.abs(param.data) < threshold] = 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5018969",
   "metadata": {},
   "source": [
    "## Represent as sparse matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "94b4e2b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.sparse as sp\n",
    "\n",
    "sparse_weights = []\n",
    "\n",
    "# Convert the pruned weights to a sparse matrix\n",
    "for name, param in net.named_parameters():\n",
    "    if 'weight' in name:\n",
    "        \n",
    "        np_weight = param.data.cpu().numpy()\n",
    "        sparse_weights.append(sp.csr_matrix(np_weight))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4637db3",
   "metadata": {},
   "source": [
    "## Size before pruning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "28196989",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2.1611328125"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "total_size = 0\n",
    "\n",
    "for name, param in net.named_parameters():\n",
    "    if 'weight' in name:\n",
    "        tensor = param.data\n",
    "        total_size += tensor.element_size() * tensor.numel()\n",
    "        \n",
    "# Convert bytes to a more human-readable format (e.g., megabytes)\n",
    "tensor_size_mb = total_size/(1024**2)\n",
    "\n",
    "tensor_size_mb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8526f7fc",
   "metadata": {},
   "source": [
    "## Size after pruning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "da788299",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.8427848815917969"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "total_size = 0\n",
    "\n",
    "for w in sparse_weights:\n",
    "    total_size += w.data.nbytes\n",
    "    \n",
    "csr_size_mb = total_size/(1024**2)\n",
    "csr_size_mb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8faa759a",
   "metadata": {},
   "source": [
    "# Activation pruning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d73fdf3",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "baa723a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a simple teacher neural network with 4 fully connected layers\n",
    "class SimpleNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(SimpleNet, self).__init__()\n",
    "        self.fc1 = nn.Linear(28 * 28, 512)\n",
    "        self.fc2 = nn.Linear(512, 256)\n",
    "        self.fc3 = nn.Linear(256, 128)\n",
    "        self.fc4 = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x1 = torch.relu(self.fc1(x))\n",
    "        x2 = torch.relu(self.fc2(x1))\n",
    "        x3 = torch.relu(self.fc3(x2))\n",
    "        x4 = self.fc4(x3)\n",
    "        return x1, x2, x3, x4  # Return intermediate feature activations for activation pruning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "6672e69c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model):\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    with torch.no_grad():\n",
    "        for data in testloader:\n",
    "            inputs, labels = data\n",
    "            outputs = model(inputs)[-1] # use last element returned by forward function\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "    return correct / total"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab7e44f0",
   "metadata": {},
   "source": [
    "## Initialize and train the neural network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "8ba3cc52",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Loss: 0.31967253410525476, Accuracy: 95.16%\n",
      "Epoch 2, Loss: 0.1467606759618229, Accuracy: 96.23%\n",
      "Epoch 3, Loss: 0.10635705036061532, Accuracy: 96.97%\n",
      "Epoch 4, Loss: 0.09333869409032547, Accuracy: 97.34%\n",
      "Epoch 5, Loss: 0.07620950450357208, Accuracy: 96.63%\n"
     ]
    }
   ],
   "source": [
    "net = SimpleNet()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(net.parameters(), lr=0.001)\n",
    "\n",
    "for epoch in range(5):\n",
    "    net.train()\n",
    "    running_loss = 0.0\n",
    "    \n",
    "    for data in trainloader:\n",
    "        inputs, labels = data\n",
    "        optimizer.zero_grad()\n",
    "        outputs = net(inputs)\n",
    "        loss = criterion(outputs[-1], labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        running_loss += loss.item()\n",
    "        \n",
    "    accuracy = evaluate(net)\n",
    "        \n",
    "    print(f\"Epoch {epoch + 1}, Loss: {running_loss / len(trainloader)}, Accuracy: {accuracy * 100:.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e88fab01",
   "metadata": {},
   "source": [
    "## Evaluate the average activations of neurons on the training data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "9fd624a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "net.eval()\n",
    "all_activations = [torch.zeros(512), torch.zeros(256), torch.zeros(128)]\n",
    "data_size = len(trainloader.dataset.targets)\n",
    "\n",
    "with torch.no_grad():\n",
    "    for data in trainloader:\n",
    "        inputs, _ = data\n",
    "        \n",
    "        activations_fc1 = torch.relu(net.fc1(inputs.view(-1, 28*28)))\n",
    "        activations_fc2 = torch.relu(net.fc2(activations_fc1))\n",
    "        activations_fc3 = torch.relu(net.fc3(activations_fc2))\n",
    "        \n",
    "        all_activations[0] += torch.sum(activations_fc1, dim=0)\n",
    "        all_activations[1] += torch.sum(activations_fc2, dim=0)\n",
    "        all_activations[2] += torch.sum(activations_fc3, dim=0)\n",
    "            \n",
    "for idx, activations in enumerate(all_activations):\n",
    "    all_activations[idx] = activations/data_size"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53dc0138",
   "metadata": {},
   "source": [
    "## Apply activation pruning across thresholds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "40db69db",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████| 11/11 [01:44<00:00,  9.54s/it]\n"
     ]
    }
   ],
   "source": [
    "thresholds = np.linspace(0, 1, 11)\n",
    "results = []\n",
    "\n",
    "original_total_params = sum(p.numel() for p in net.parameters())\n",
    "for threshold in tqdm(thresholds):\n",
    "    new_net = SimpleNet()\n",
    "    \n",
    "    new_net.fc1.weight = net.fc1.weight\n",
    "    new_net.fc2.weight = net.fc2.weight\n",
    "    new_net.fc3.weight = net.fc3.weight\n",
    "    new_net.fc4.weight = net.fc4.weight\n",
    "    \n",
    "    new_net.fc1.bias = net.fc1.bias\n",
    "    new_net.fc2.bias = net.fc2.bias\n",
    "    new_net.fc3.bias = net.fc3.bias\n",
    "    new_net.fc4.bias = net.fc4.bias\n",
    "    \n",
    "    new_net.fc1.weight = nn.Parameter(new_net.fc1.weight[all_activations[0]>=threshold])\n",
    "\n",
    "    new_net.fc2.weight = nn.Parameter(new_net.fc2.weight[:, all_activations[0]>=threshold])\n",
    "    new_net.fc2.weight = nn.Parameter(new_net.fc2.weight[all_activations[1]>=threshold])\n",
    "    \n",
    "    new_net.fc3.weight = nn.Parameter(new_net.fc3.weight[:, all_activations[1]>=threshold])\n",
    "    new_net.fc3.weight = nn.Parameter(new_net.fc3.weight[all_activations[2]>=threshold])\n",
    "    \n",
    "    new_net.fc4.weight = nn.Parameter(new_net.fc4.weight[:, all_activations[2]>=threshold])\n",
    "    \n",
    "    \n",
    "    new_net.fc1.bias = nn.Parameter(new_net.fc1.bias[all_activations[0]>=threshold])\n",
    "    new_net.fc2.bias = nn.Parameter(new_net.fc2.bias[all_activations[1]>=threshold])\n",
    "    new_net.fc3.bias = nn.Parameter(new_net.fc3.bias[all_activations[2]>=threshold])\n",
    "    \n",
    "    \n",
    "    accuracies = 0\n",
    "    total_time = 0\n",
    "    for _ in range(7):\n",
    "        \n",
    "        start = time()\n",
    "        accuracies += evaluate(new_net)    \n",
    "        total_time += time() - start\n",
    "        \n",
    "    new_total_params = sum(p.numel() for p in new_net.parameters())\n",
    "\n",
    "    results.append([threshold, 100*accuracies/7, original_total_params, new_total_params, total_time/7])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "2b2adf70",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><div id=eedf4b62-be05-4fd5-89b9-6fe8c623b6d1 style=\"display:none; background-color:#9D6CFF; color:white; width:200px; height:30px; padding-left:5px; border-radius:4px; flex-direction:row; justify-content:space-around; align-items:center;\" onmouseover=\"this.style.backgroundColor='#BA9BF8'\" onmouseout=\"this.style.backgroundColor='#9D6CFF'\" onclick=\"window.commands?.execute('create-mitosheet-from-dataframe-output');\">See Full Dataframe in Mito</div> <script> if (window.commands?.hasCommand('create-mitosheet-from-dataframe-output')) document.getElementById('eedf4b62-be05-4fd5-89b9-6fe8c623b6d1').style.display = 'flex' </script> <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Threshold</th>\n",
       "      <th>Accuracy</th>\n",
       "      <th>Original Params</th>\n",
       "      <th>New Params</th>\n",
       "      <th>Inference Time</th>\n",
       "      <th>Size Reduction</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>0.0</td>\n",
       "      <td>96.63</td>\n",
       "      <td>567434</td>\n",
       "      <td>567434</td>\n",
       "      <td>1.405958</td>\n",
       "      <td>0.000000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>0.1</td>\n",
       "      <td>96.44</td>\n",
       "      <td>567434</td>\n",
       "      <td>215901</td>\n",
       "      <td>1.196406</td>\n",
       "      <td>0.619513</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>0.2</td>\n",
       "      <td>96.37</td>\n",
       "      <td>567434</td>\n",
       "      <td>199489</td>\n",
       "      <td>1.361184</td>\n",
       "      <td>0.648437</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>0.3</td>\n",
       "      <td>96.11</td>\n",
       "      <td>567434</td>\n",
       "      <td>182201</td>\n",
       "      <td>1.425039</td>\n",
       "      <td>0.678904</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>0.4</td>\n",
       "      <td>94.24</td>\n",
       "      <td>567434</td>\n",
       "      <td>165041</td>\n",
       "      <td>1.300359</td>\n",
       "      <td>0.709145</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>0.6</td>\n",
       "      <td>87.49</td>\n",
       "      <td>567434</td>\n",
       "      <td>136019</td>\n",
       "      <td>1.325743</td>\n",
       "      <td>0.760291</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>0.7</td>\n",
       "      <td>86.40</td>\n",
       "      <td>567434</td>\n",
       "      <td>124319</td>\n",
       "      <td>1.423424</td>\n",
       "      <td>0.780910</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>0.8</td>\n",
       "      <td>82.55</td>\n",
       "      <td>567434</td>\n",
       "      <td>106310</td>\n",
       "      <td>1.495066</td>\n",
       "      <td>0.812648</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>0.9</td>\n",
       "      <td>77.36</td>\n",
       "      <td>567434</td>\n",
       "      <td>94773</td>\n",
       "      <td>1.370770</td>\n",
       "      <td>0.832980</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>1.0</td>\n",
       "      <td>64.49</td>\n",
       "      <td>567434</td>\n",
       "      <td>76409</td>\n",
       "      <td>1.262326</td>\n",
       "      <td>0.865343</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table></div>"
      ],
      "text/plain": [
       "    Threshold  Accuracy  Original Params  New Params  Inference Time  \\\n",
       "0         0.0     96.63           567434      567434        1.405958   \n",
       "1         0.1     96.44           567434      215901        1.196406   \n",
       "2         0.2     96.37           567434      199489        1.361184   \n",
       "3         0.3     96.11           567434      182201        1.425039   \n",
       "4         0.4     94.24           567434      165041        1.300359   \n",
       "5         0.5     93.28           567434      146879        1.400834   \n",
       "6         0.6     87.49           567434      136019        1.325743   \n",
       "7         0.7     86.40           567434      124319        1.423424   \n",
       "8         0.8     82.55           567434      106310        1.495066   \n",
       "9         0.9     77.36           567434       94773        1.370770   \n",
       "10        1.0     64.49           567434       76409        1.262326   \n",
       "\n",
       "    Size Reduction  \n",
       "0         0.000000  \n",
       "1         0.619513  \n",
       "2         0.648437  \n",
       "3         0.678904  \n",
       "4         0.709145  \n",
       "5         0.741152  \n",
       "6         0.760291  \n",
       "7         0.780910  \n",
       "8         0.812648  \n",
       "9         0.832980  \n",
       "10        0.865343  "
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results = pd.DataFrame(results, columns = [\"Threshold\", \"Accuracy\", \"Original Params\", \"New Params\", \"Inference Time\"])\n",
    "results[\"Size Reduction\"] = 1-results[\"New Params\"]/results[\"Original Params\"]\n",
    "results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee6f38ae",
   "metadata": {},
   "source": [
    "# Low-rank Factorization"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c914d5b",
   "metadata": {},
   "source": [
    "## Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "99b2534c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a simple teacher neural network with 4 fully connected layers\n",
    "class SimpleNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(SimpleNet, self).__init__()\n",
    "        self.fc1 = nn.Linear(28 * 28, 512)\n",
    "        self.fc2 = nn.Linear(512, 256)\n",
    "        self.fc3 = nn.Linear(256, 128)\n",
    "        self.fc4 = nn.Linear(128, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.view(x.size(0), -1)\n",
    "        x1 = torch.relu(self.fc1(x))\n",
    "        x2 = torch.relu(self.fc2(x1))\n",
    "        x3 = torch.relu(self.fc3(x2))\n",
    "        x4 = self.fc4(x3)\n",
    "        return x4 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "b1929bf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model):\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    total = 0\n",
    "    with torch.no_grad():\n",
    "        for data in testloader:\n",
    "            inputs, labels = data\n",
    "            outputs = model(inputs) # use last element returned by forward function\n",
    "            _, predicted = torch.max(outputs.data, 1)\n",
    "            total += labels.size(0)\n",
    "            correct += (predicted == labels).sum().item()\n",
    "    return correct / total"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "275d8edb",
   "metadata": {},
   "source": [
    "## Initialize and train the neural network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "50687b4d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1, Loss: 0.32092429239596765, Accuracy: 94.65%\n",
      "Epoch 2, Loss: 0.13917262043887332, Accuracy: 96.70%\n",
      "Epoch 3, Loss: 0.10783804802925014, Accuracy: 96.63%\n",
      "Epoch 4, Loss: 0.08807487449378394, Accuracy: 96.81%\n",
      "Epoch 5, Loss: 0.07418891476260732, Accuracy: 97.24%\n"
     ]
    }
   ],
   "source": [
    "net = SimpleNet()\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "optimizer = optim.Adam(net.parameters(), lr=0.001)\n",
    "\n",
    "for epoch in range(5):\n",
    "    net.train()\n",
    "    running_loss = 0.0\n",
    "    \n",
    "    for data in trainloader:\n",
    "        inputs, labels = data\n",
    "        optimizer.zero_grad()\n",
    "        outputs = net(inputs)\n",
    "        loss = criterion(outputs, labels)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        \n",
    "        running_loss += loss.item()\n",
    "        \n",
    "    accuracy = evaluate(net)\n",
    "        \n",
    "    print(f\"Epoch {epoch + 1}, Loss: {running_loss / len(trainloader)}, Accuracy: {accuracy * 100:.2f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c13c389b",
   "metadata": {},
   "source": [
    "## Function to determine minimum matrix operations\n",
    "\n",
    "Taken from https://www.geeksforgeeks.org/python-program-for-matrix-chain-multiplication-dp-8/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "07996bbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "def MatrixChainOrder(p, i, j):\n",
    " \n",
    "    if i == j:\n",
    "        return 0\n",
    " \n",
    "    _min = sys.maxsize\n",
    "     \n",
    "    for k in range(i, j):\n",
    "     \n",
    "        count = (MatrixChainOrder(p, i, k)\n",
    "             + MatrixChainOrder(p, k + 1, j)\n",
    "                   + p[i-1] * p[k] * p[j])\n",
    " \n",
    "        if count < _min:\n",
    "            _min = count;\n",
    "     \n",
    "    return _min;"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a9bec39",
   "metadata": {},
   "source": [
    "## Apply LRF across thresholds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "b6174ee7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████| 13/13 [02:16<00:00, 10.47s/it]\n"
     ]
    }
   ],
   "source": [
    "rank_values = [128, 100, 90, 80, 60, 50, 40, 30, 20, 10, 5, 2, 1]\n",
    "results = []\n",
    "original_total_params = 128*256\n",
    "batch_size = 32\n",
    "\n",
    "# factorize layer 3 weights\n",
    "U, S, V = torch.svd(net.fc3.weight)\n",
    "\n",
    "for rank in tqdm(rank_values):\n",
    "\n",
    "    # Truncate U, S, and V to retain only the top 'rank' components\n",
    "    U_low_rank = U[:, :rank]\n",
    "    S_low_rank = torch.diag(S[:rank])\n",
    "    V_low_rank = V[:, :rank]\n",
    "\n",
    "    # Reconstruct the factorized weight matrix\n",
    "    factorized_weight_matrix = torch.mm(U_low_rank, torch.mm(S_low_rank, V_low_rank.t()))\n",
    "\n",
    "    # Replace the weight matrix of the chosen layer with the factorized weight matrix\n",
    "    net.fc3.weight = nn.Parameter(factorized_weight_matrix)\n",
    "    \n",
    "    weight_list = [batch_size, 256, rank, rank, 128]\n",
    "    \n",
    "    if rank == 128:\n",
    "        # use the usual weight matrix\n",
    "        total_operations = batch_size*256*128\n",
    "    else:\n",
    "        # find the minimum number of operations\n",
    "        total_operations = MatrixChainOrder(weight_list, 1, 4)\n",
    "\n",
    "    accuracies = 0\n",
    "    total_time = 0\n",
    "    for _ in range(7):\n",
    "        \n",
    "        start = time()\n",
    "        accuracies += evaluate(net)    \n",
    "        total_time += time() - start\n",
    "    \n",
    "    # total parameters of three matrices\n",
    "    new_total_params = 128*rank + rank**2 + rank*256\n",
    "\n",
    "    results.append([rank, 100*accuracies/7, original_total_params, new_total_params, total_operations, total_time/7])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "1d9937e1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><div id=5762aa58-7b57-46bc-bdce-d71b97e96ea3 style=\"display:none; background-color:#9D6CFF; color:white; width:200px; height:30px; padding-left:5px; border-radius:4px; flex-direction:row; justify-content:space-around; align-items:center;\" onmouseover=\"this.style.backgroundColor='#BA9BF8'\" onmouseout=\"this.style.backgroundColor='#9D6CFF'\" onclick=\"window.commands?.execute('create-mitosheet-from-dataframe-output');\">See Full Dataframe in Mito</div> <script> if (window.commands?.hasCommand('create-mitosheet-from-dataframe-output')) document.getElementById('5762aa58-7b57-46bc-bdce-d71b97e96ea3').style.display = 'flex' </script> <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Threshold</th>\n",
       "      <th>Accuracy</th>\n",
       "      <th>Original Params</th>\n",
       "      <th>New Params</th>\n",
       "      <th>Operations</th>\n",
       "      <th>Inference Time</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>128</td>\n",
       "      <td>97.24</td>\n",
       "      <td>32768</td>\n",
       "      <td>65536</td>\n",
       "      <td>1048576</td>\n",
       "      <td>1.272552</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>100</td>\n",
       "      <td>97.24</td>\n",
       "      <td>32768</td>\n",
       "      <td>48400</td>\n",
       "      <td>1548800</td>\n",
       "      <td>1.350363</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>90</td>\n",
       "      <td>97.26</td>\n",
       "      <td>32768</td>\n",
       "      <td>42660</td>\n",
       "      <td>1365120</td>\n",
       "      <td>1.414365</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>80</td>\n",
       "      <td>97.26</td>\n",
       "      <td>32768</td>\n",
       "      <td>37120</td>\n",
       "      <td>1187840</td>\n",
       "      <td>1.439593</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>60</td>\n",
       "      <td>97.26</td>\n",
       "      <td>32768</td>\n",
       "      <td>26640</td>\n",
       "      <td>852480</td>\n",
       "      <td>1.411347</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>20</td>\n",
       "      <td>97.16</td>\n",
       "      <td>32768</td>\n",
       "      <td>8080</td>\n",
       "      <td>258560</td>\n",
       "      <td>1.571930</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>10</td>\n",
       "      <td>97.24</td>\n",
       "      <td>32768</td>\n",
       "      <td>3940</td>\n",
       "      <td>126080</td>\n",
       "      <td>1.524304</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>5</td>\n",
       "      <td>80.23</td>\n",
       "      <td>32768</td>\n",
       "      <td>1945</td>\n",
       "      <td>62240</td>\n",
       "      <td>1.817376</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>2</td>\n",
       "      <td>34.73</td>\n",
       "      <td>32768</td>\n",
       "      <td>772</td>\n",
       "      <td>24704</td>\n",
       "      <td>1.589601</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>1</td>\n",
       "      <td>19.08</td>\n",
       "      <td>32768</td>\n",
       "      <td>385</td>\n",
       "      <td>12320</td>\n",
       "      <td>1.575846</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table></div>"
      ],
      "text/plain": [
       "    Threshold  Accuracy  Original Params  New Params  Operations  \\\n",
       "0         128     97.24            32768       65536     1048576   \n",
       "1         100     97.24            32768       48400     1548800   \n",
       "2          90     97.26            32768       42660     1365120   \n",
       "3          80     97.26            32768       37120     1187840   \n",
       "4          60     97.26            32768       26640      852480   \n",
       "5          50     97.23            32768       21700      694400   \n",
       "6          40     97.25            32768       16960      542720   \n",
       "7          30     97.23            32768       12420      397440   \n",
       "8          20     97.16            32768        8080      258560   \n",
       "9          10     97.24            32768        3940      126080   \n",
       "10          5     80.23            32768        1945       62240   \n",
       "11          2     34.73            32768         772       24704   \n",
       "12          1     19.08            32768         385       12320   \n",
       "\n",
       "    Inference Time  \n",
       "0         1.272552  \n",
       "1         1.350363  \n",
       "2         1.414365  \n",
       "3         1.439593  \n",
       "4         1.411347  \n",
       "5         1.425996  \n",
       "6         1.498162  \n",
       "7         1.538632  \n",
       "8         1.571930  \n",
       "9         1.524304  \n",
       "10        1.817376  \n",
       "11        1.589601  \n",
       "12        1.575846  "
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "results = pd.DataFrame(results, columns = [\"Threshold\", \"Accuracy\", \"Original Params\", \"New Params\", \"Operations\", \"Inference Time\"])\n",
    "results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7efefc9",
   "metadata": {},
   "source": [
    "# Quantization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "79866360",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Quantize the model to int8 data type\n",
    "quantized_model = torch.quantization.quantize_dynamic(\n",
    "    net,  # Your model\n",
    "    {torch.nn.Linear},  # Specify the layers to be quantized (e.g., linear layers)\n",
    "    dtype=torch.qint8  # Data type for quantization (int8)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "1f62e19e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "97.24%\n"
     ]
    }
   ],
   "source": [
    "evaluate(quantized_model)"
   ]
  }
 ],
 "metadata": {
  "finalized": {
   "timestamp": 1694847489806,
   "trusted": false
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
